# Copyright Tashan and Friends. 2010

from pylab import zeros_like,size,sin,plot,show,zeros,ones_like,legend,xlabel,ylabel,grid,linspace,array,title


has_pbar = True
try:
    import progressbar
except ImportError:
    has_pbar = False
    
class PBar(object):
    """A simple wrapper around the progressbar so it works if a user has
    it installed or not.
    """
    def __init__(self, maxval, show=True):
        bar = None
        self.count = 0
        if has_pbar and show:
            widgets = progressbar.default_widgets + \
                      [progressbar.ETA()]
            bar = progressbar.ProgressBar(widgets=widgets,
                                          maxval=maxval).start()
        self.bar = bar

    def update(self):
        if self.bar is not None:
            self.bar.update(self.count)
            self.count += 1

    def finish(self):
        if self.bar is not None:
            self.bar.finish()




class StateSpaceSolver:
    """This is the base class for solving any state space problem.
       Input:
       In<type="class"> - Is the class of input given to the system, currently available options are:
       1. Step
       2. Ramp
       3. Impulse
       4. Sin
       5. Poly
       type "Step?" for more info about each class
       
       dcoeffs<type="array"> - Coefficients of characteristic equation in ascending order of power of s.
       ncoeffs<type="array"> - Coefficients of numerator in ascending order of power of s.
       dt<type="double"> - Time step of the simulation, default value is set to 1e-2
    
       Example: Suppose the trasfer function is given by num/den where den = 3s^2 + 2s + 1 and num = 1 and input is step of magnitude 2.0. It has to be solve for time 5 sec with dt=1e-3. Then:
       In = Step(2.0)
       den = array([1.0,2.0,3.0])
       num = array([1.0])
       a = StateSpaceSolver(In, den, num, 1e-3)
       a.solver(5)
       
       It will automatically plot the Input and Output
       Input, Output and Time can be access by using a.Inp, a.Out, and a.Tim respectively.
    """

    def __init__(self, In,  dcoeffs, ncoeffs, dt = 1e-2):
        self.In = In
        self.dcoeffs = dcoeffs
        self.ncoeffs = ncoeffs
        self.dt = dt
        self.p = size(self.dcoeffs)
        self.x = zeros([self.p-1])
        self.xdot = zeros([self.p-1])
        self.n = self.p
        self.re = 0.0
        self.nn = size(self.ncoeffs)

    def _solve(self, ti):
        dt = self.dt
        n = self.n
        x = self.x
        dcoeffs = self.dcoeffs
        xdot = self.xdot
        xdot[n-2] = self.In.eval(ti)
        self.tmp3.append(xdot[n-2])
        for i in range(n-1): # max value of i will be n-2
            xdot[n-2] -= dcoeffs[i]*x[i] 
      
        for i in range(3,n+1): # i going from 2 to n-1 means j going from n-2 to 1
            j = n-i
            x[j+1] += xdot[j+1]*dt
            xdot[j] = x[j+1]
            
   
        x[0] += xdot[0]*dt
        self.re = 0.0
        for i in range(self.nn): # max value of i will be n-2
            self.re += self.ncoeffs[i]*x[i] 

        self.x = x


    def solver(self, t = 1.0):
        ti = 0.0
        self.tmp1 = []
        self.tmp2 = []
        self.tmp3 = []
        maxval = int((t - ti)/self.dt +1)
        bar = PBar(maxval, show=True)
        while ti < t:
            self._solve(ti)
            self.tmp1.append(self.re)
            self.tmp2.append(ti)
            bar.update()
            ti += self.dt
            
        plot(self.tmp2, self.tmp1)
        plot(self.tmp2, self.tmp3)
        
        self.Inp = self.tmp3
        self.Out = self.tmp1
        self.Tim = self.tmp2

        legend(['Output','Input'])
        xlabel('Time(s)')
        ylabel('Output')
        
        tmps=''
        for i in self.dcoeffs:
            tmps = tmps + '%f' %i
        print tmps
        
        
        grid(b=None, which='major')
        show()
            
def frequency_response():
    """This doesnt work"""
    a = linspace(0,1000,1000)
    tmp = []
    for i in a:
        var = StateSpaceSolver(sin,array([1,1,1]),array([1]), f=i)
        var.solver(10)
        tmp.append(max(var.tmp1))
    return tmp





#########################################################################
#                                                                       #
#                          Declaring Input Class                        #
#                                                                       #
#########################################################################

class Step:
    """
    Step Function
    Input:A<type="double"> - Magnitude of the step
    """
    def __init__(self, A = 1.0):
        self.A = A
    def eval(self, x):
        return self.A

class Ramp:
    """
    Ramp Function
    Input:m<type="double"> - Slope of the ramp
    """
    def __init__(self, m = 1.0):
        self.m = m
    def eval(self, x):
        return self.m*x

class Impulse:
    """
    Impulse Function
    Input: A<type="double"> - Area under the impulse function
    dt<type="double"> - Time for which impulse act
    """
    def __init__(self, A = 1.0, dt = 1e-2):
        self.mag = A/dt
        self.dt = dt
    def eval(self, x):
        if x<self.dt:
            return self.mag
        else:
            return 0.0

class Sin:
    """
    Sine Function
    Input: A<type="double"> - Magnitude of wave 
    f<type="double"> - frequency of sine wave
    phi<type="double"> - initial phase angle, by default set to 0
    """
    def __init__(self, A =1.0, f = 1.0, phi = 0.0):
        self.f = f
        self.phi = phi
        self.A = A
    def eval(self, x):
        return self.A*sin(self.f*x + self.phi)

class Poly:
    """
    Polynomial Function, generating polynomial from the give coeffs
    Input: coeffs<type="list"> - List of coefficient of polynomial in ascending order of power
    """
    def __init__(self, coeffs = [1]):
        self.n = size(coeffs)
        self.coeffs = coeffs
    def eval(self, x):
        tmp = 0
        for i in range(self.n):
            tmp += self.coeffs[i]*x**i
        return tmp


